import math
from collections import Counter, defaultdict


def calculate_tf(words):
    # 统计每个词在本文档中出现的次数
    word_count = Counter(words)  # 人工智能:1.....
    # 统计总词数
    total_words = len(words)  # 5
    return {word: count / total_words for word, count in word_count.items()}


def calculate_idf(documents):
    # 统计文档的总数
    total_docs = len(documents)  # 5
    # 统计每个词出现在多少个文档中
    word_doc_count = defaultdict(int)
    # 遍历文档
    for idx, doc in enumerate(documents):
        # 获取文档中的唯一词集合 如果一个词在同一个文档中出现了多次，只能算一次
        unique_words = set(doc)
        # 统计每个词出现的文档数
        for word in unique_words:
            word_doc_count[word] += 1  # word_doc_count["人工智能"]=4
    res = {}
    for word, count in word_doc_count.items():
        inversedocumentFrequency = total_docs / (count + 1)
        value = math.log(inversedocumentFrequency)
        res[word] = value
    return res


def calculate_tfidf(documents):
    # 计算每个文档的TF值
    tf_scores = []
    for idx, doc in enumerate(documents):
        tf = calculate_tf(doc)
        tf_scores.append(tf)

    # 计算IDF值
    idf_scores = calculate_idf(documents)

    tfidf_scores = []
    for idx, tf_doc in enumerate(tf_scores):
        doc_tfidf = {}
        for word, tf in tf_doc.items():
            doc_tfidf[word] = tf * idf_scores[word]
        tfidf_scores.append(doc_tfidf)
    return tfidf_scores


if __name__ == "__main__":
    # 文档
    documents = [
        ["人工智能", "计算机科学", "分支", "创建", "系统"],
        ["机器学习", "人工智能", "子领域", "计算机", "学习"],
        ["深度学习", "机器学习", "分支", "神经网络", "学习"],
        ["自然语言处理", "人工智能", "应用", "计算机", "语言"],
        ["计算机视觉", "人工智能", "分支", "计算机", "视觉"],
    ]
    tfidf_results = calculate_tfidf(documents)
    for i, doc_tfidf in enumerate(tfidf_results):
        print(f"文档{i+1}的TF-IDF值:")
        sorted_words = sorted(doc_tfidf.items(), key=lambda x: x[1], reverse=True)
        for word, score in sorted_words:
            print(f"{word}:{score:.4f}")
        print()
